package com.sunzm.spark.sql.udaf

import java.io.{ByteArrayInputStream, ByteArrayOutputStream, DataInputStream, DataOutputStream}

import org.roaringbitmap.RoaringBitmap

/**
 * BitMap操作工具
 */
object BitMapUtil {

  /**
   * 把一个 RoaringBitmap 序列化成byte数组
   * @param bm
   * @return
   */
  def serBitMap(bm:RoaringBitmap):Array[Byte] = {
    val stream = new ByteArrayOutputStream()
    val dataOutput = new DataOutputStream(stream)

    bm.serialize(dataOutput)

    stream.toByteArray
  }


  /**
   * 把一个 byte数组 反序列化为一个 RoaringBitmap
   * @param byteArray
   * @return
   */
  def deSerBitMap(byteArray:Array[Byte]):RoaringBitmap = {

    val bitmap: RoaringBitmap = RoaringBitmap.bitmapOf()

    val stream = new ByteArrayInputStream(byteArray)
    val dataInput = new DataInputStream(stream)

    bitmap.deserialize(dataInput)

    bitmap
  }

  /**
   * 合并两个序列化后的Bitmap, 并把就结果序列化后返回
   * @param buffer1
   * @param buffer2
   * @return
   */
  def mergeBitmapByteArray(buffer1: Array[Byte], buffer2: Array[Byte]): Array[Byte] = {

    //反序列化buffer1
    val bitmap1: RoaringBitmap = BitMapUtil.deSerBitMap(buffer1)
    //反序列化buffer2
    val bitmap2: RoaringBitmap = BitMapUtil.deSerBitMap(buffer2)

    // 合并两个bitmap（或操作）,  最终的结果会更新 bitmap1
    bitmap1.or(bitmap2)

    //将合并后的 bitmap 再序列化成字节数组
    BitMapUtil.serBitMap(bitmap1)
  }
}
